# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #
# ~~~~~~~ 利用DCT隐写术对图像中的秘密信息进行编码 ~~~~~~~ #
# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #

from skimage.io import imread, imsave
from skimage.color import rgb2ycbcr, ycbcr2rgb
from skimage import img_as_float, img_as_ubyte
import numpy as np
import sys

from DataHidingDemo.dct_processing import dct_process_channel as dct
from DataHidingDemo.dct_processing import idct_process_channel as idct
from DataHidingDemo.extra import file_to_bin, bin_to_file, mes_to_bin


# 库无法读取示例中的一些图像
from PIL import PngImagePlugin

PngImagePlugin.MAX_TEXT_CHUNK = 500 * (1024 ** 2)

# 该值位于区间[1；99]中，
# 这决定了图像质量
_quality = 70


# TODO: 更多的文件格式支持
# TODO: 根据DCT特定中频系数大小关系来隐写数据


# 输入：RGB图像 ，带扩展名的文件名（secret）
# 输出：保存PNG stego图像
def dct_hide(img, mes="", filename="", out_name='output', N=8):
    # 保存原始图像
    original = img.copy()

    # 找出图像的大小和频道数
    a, b = img.shape[:2]
    c = 3 if len(img.shape) == 3 else 1


    if mes == "":
        # 接收的秘密消息转化比特流
        data = file_to_bin(filename)

    else:
        # 接收的秘密消息转化比特流
        data= mes_to_bin(mes)

    # 检查图像的容量
    if len(data) > (a // 8) * (b // 8) * 3:
        print("File is too large")
        return

    index = 0
    if c == 3:
        newimg = rgb2ycbcr(img)

        # 计算每个信道上的DCT系数
        Y = dct(newimg[..., 0], _quality)
        Cb = dct(newimg[..., 1], _quality)
        Cr = dct(newimg[..., 2], _quality)
        newimg = np.dstack((Y, Cb, Cr))

        # 将LSB方法应用于每个NxN块中的系数
        for i in range(0, a, 8):
            for j in range(0, b, 8):
                for k in range(c):
                    if index + 1 < len(data):
                        newimg[i + 1, j + 1, k] += -(int(newimg[i + 1, j + 1, k]) % 2) + int(data[index])
                        index += 1

        # 逆DCT以获得新的YCbCr像素值
        Y = idct(newimg[..., 0], _quality)
        Cb = idct(newimg[..., 1], _quality)
        Cr = idct(newimg[..., 2], _quality)
        newimg = np.dstack((Y, Cb, Cr))

        # 转换到RGB颜色空间
        img = np.clip(0, 1.0, ycbcr2rgb(newimg))
        img = img_as_ubyte(img)
        return img,original


def main():
    print("Enter the source path of the image to encode:")
    file = input()
    img = imread(file)
    print("Enter message to be hidden: ")
    file2 = input()
    dct_hide(img, file2, 'output')


if __name__ == '__main__':
    main()
